import datetime
import os
import subprocess
import wave

import pyaudio


def record_system_audio(filename, duration=5, sample_rate=48000, channels=2, chunk=1024):
    """
    录制系统输出音频（扬声器播放的声音）
    :param filename: 输出文件名
    :param duration: 录制时长(秒)
    """
    p = pyaudio.PyAudio()

    # 查找立体声混音设备
    device_index = None
    for i in range(p.get_device_count()):
        dev = p.get_device_info_by_index(i)
        if "立体声混音" in dev["name"] or "Stereo Mix" in dev["name"]:
            device_index = dev["index"]
            print(f"找到立体声混音设备: {dev['name']}")
            break

    if device_index is None:
        raise Exception("未找到立体声混音设备，请检查Windows音频设置")

    stream = p.open(
        format=pyaudio.paInt16,
        channels=channels,
        rate=sample_rate,
        input=True,
        input_device_index=device_index,
        frames_per_buffer=chunk
    )

    print("● 开始录制系统音频...")
    frames = []

    for _ in range(0, int(sample_rate / chunk * duration)):
        frames.append(stream.read(chunk))

    print("■ 录制完成")

    stream.stop_stream()
    stream.close()
    p.terminate()

    # 保存WAV文件
    with wave.open(filename, 'wb') as wf:
        wf.setnchannels(channels)
        wf.setsampwidth(p.get_sample_size(pyaudio.paInt16))
        wf.setframerate(sample_rate)
        wf.writeframes(b''.join(frames))

    print(f"✓ 系统音频已保存: {filename}")


def convert_wav_to_mp3_ffmpeg(input_wav, output_mp3, bitrate="320k"):
    """
    使用原生FFmpeg转换WAV到MP3
    :param input_wav: 输入WAV文件路径
    :param output_mp3: 输出MP3文件路径
    :param bitrate: MP3比特率 (128k, 192k, 256k, 320k)
    """
    cmd = [
        "ffmpeg",
        "-i", input_wav,
        "-vn",  # 忽略视频流
        "-acodec", "libmp3lame",  # 使用MP3编码器
        "-ab", bitrate,  # 音频比特率
        "-ar", "44100",  # 采样率
        "-y",  # 覆盖输出文件
        output_mp3
    ]

    try:
        result = subprocess.run(cmd, check=True, capture_output=True, text=True)
        if result.returncode == 0:
            print(f"✓ 转换成功: {input_wav} → {output_mp3}")
            print(f"  文件大小: {os.path.getsize(output_mp3) // 1024}KB")
        return True
    except subprocess.CalledProcessError as e:
        print(f"× 转换失败: {e.stderr}")
        return False
    except FileNotFoundError:
        print("× FFmpeg未找到，请确保已正确安装FFmpeg")
        return False


if __name__ == "__main__":
    date_string = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
    # wav 文件
    wav_filename = "E:/test/" + date_string + "-system_audio.wav"
    # mp3 文件
    mp3_filename = "E:/test/" + date_string + "-system_audio.mp3"
    # 录制时长（秒）
    duration_time = 180

    # 录制系统输出音频（扬声器播放的声音）
    try:
        record_system_audio(wav_filename, duration=duration_time)
    except Exception as e:
        print(f"错误: {str(e)}")
        print("请执行以下步骤:")
        print("1. 右键点击任务栏音量图标 -> 打开声音设置")
        print("2. 选择'声音控制面板'（右侧相关设置）")
        print("3. 在'录制'选项卡中右键 -> 显示禁用的设备")
        print("4. 启用'立体声混音'设备")
        print("5. 将其设置为默认设备")

    # convert wav to mp3
    convert_wav_to_mp3_ffmpeg(
        wav_filename,
        mp3_filename,
        bitrate="320k")
